# -*- coding: utf-8 -*-
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Circle, Rectangle
import matplotlib.animation as animation
from matplotlib.collections import PatchCollection
import math

# 参数设置
SAFE_DISTANCE = 1.5  # 安全距离
GRID_SIZE = 30       # 网格大小
MAX_ITERATIONS = 20  # 最大迭代次数

def setup_scenario():
    # 创建无人机 - 更密集的布局以增加交互
    drones = [
        Drone(15.0, 15.0, 1),
        Drone(16.2, 16.8, 2),
        Drone(17.5, 15.3, 3),
        Drone(16.8, 14.1, 4),
        Drone(15.6, 13.9, 5)
    ]
    
    # 创建障碍物
    obstacles = [
        Obstacle(10, 15, 2, 'rectangle'),
        Obstacle(10, 20, 1, 'rectangle'),
        Obstacle(23, 18, 3, 'rectangle')
    ]
    
    return drones, obstacles

class Obstacle:
    def __init__(self, x, y, size, shape='circle'):
        self.x = x
        self.y = y
        self.size = size
        self.shape = shape
        
    def distance_to(self, x, y):
        if self.shape == 'circle':
            return max(0, math.sqrt((x - self.x)**2 + (y - self.y)**2) - self.size)
        else:  # rectangle
            dx = max(self.x - self.size - x, x - self.x - self.size, 0)
            dy = max(self.y - self.size - y, y - self.y - self.size, 0)
            return math.sqrt(dx**2 + dy**2)
            
    def check_collision(self, x, y):
        if self.shape == 'circle':
            return math.sqrt((x - self.x)**2 + (y - self.y)**2) <= self.size
        else:  # rectangle
            return (abs(x - self.x) <= self.size and abs(y - self.y) <= self.size)
            
    def check_line_collision(self, x1, y1, x2, y2):
        """检查线段是否与障碍物相交"""
        if self.shape == 'circle':
            # 计算线段到圆心的最短距离
            a = (x2 - x1)**2 + (y2 - y1)**2
            if a == 0:
                return self.check_collision(x1, y1)
                
            b = 2 * ((x2 - x1) * (x1 - self.x) + (y2 - y1) * (y1 - self.y))
            c = self.x**2 + self.y**2 + x1**2 + y1**2 - 2*(self.x*x1 + self.y*y1) - self.size**2
            
            discriminant = b**2 - 4*a*c
            if discriminant < 0:
                return False
                
            t1 = (-b + math.sqrt(discriminant)) / (2*a)
            t2 = (-b - math.sqrt(discriminant)) / (2*a)
            
            return (0 <= t1 <= 1) or (0 <= t2 <= 1)
        else:  # rectangle
            # 简化处理：检查线段是否与矩形边界相交
            rect_left = self.x - self.size
            rect_right = self.x + self.size
            rect_bottom = self.y - self.size
            rect_top = self.y + self.size
            
            # 检查线段是否与矩形的四条边相交
            return (self.line_intersects_line(x1, y1, x2, y2, rect_left, rect_bottom, rect_left, rect_top) or
                    self.line_intersects_line(x1, y1, x2, y2, rect_left, rect_top, rect_right, rect_top) or
                    self.line_intersects_line(x1, y1, x2, y2, rect_right, rect_top, rect_right, rect_bottom) or
                    self.line_intersects_line(x1, y1, x2, y2, rect_right, rect_bottom, rect_left, rect_bottom))
    
    def line_intersects_line(self, x1, y1, x2, y2, x3, y3, x4, y4):
        """检查两条线段是否相交"""
        def ccw(ax, ay, bx, by, cx, cy):
            return (cy - ay) * (bx - ax) > (by - ay) * (cx - ax)
            
        return (ccw(x1, y1, x3, y3, x4, y4) != ccw(x2, y2, x3, y3, x4, y4) and 
                ccw(x1, y1, x2, y2, x3, y3) != ccw(x1, y1, x2, y2, x4, y4))

class Drone:
    def __init__(self, x, y, id):
        self.x = x
        self.y = y
        self.id = id
        self.selected = False
        self.escape_direction = None
        self.escaped = False
        self.original_x = x
        self.original_y = y
        
    def distance_to(self, other):
        return math.sqrt((self.x - other.x)**2 + (self.y - other.y)**2)
    
    

def find_outermost_drone(drones):
    """
    找出无人机集群中最外面的点
    参数: drones: 无人机列表，每个无人机是一个元组 (x, y, id)
    返回: 最外面的无人机元组 (x, y, id)
    """
    if not drones:
        raise ValueError("无人机列表不能为空")
    center_x = sum(drone.x for drone in drones) / len(drones)
    center_y = sum(drone.y for drone in drones) / len(drones)
    max_distance = -1
    outermost_drone = None
    for drone in drones:
        x=drone.x
        y=drone.y
        drone_id=drone.id
        distance = math.sqrt((x - center_x)**2 + (y - center_y)**2)
        if distance > max_distance:
            max_distance = distance
            outermost_drone = drone
    return outermost_drone

def calculate_avoidance_direction(escape_drone, other_drones, obstacles, step_size=1.0, k_drone=1.0, k_obstacle=1.0):
    """
    计算逃离无人机的最佳逃离方向
    参数:
        escape_drone: 元组 (x, y, id) 逃离无人机的位置
        other_drones: 列表 of tuples (x, y, id) 其他无人机的位置
        obstacles: 列表 of tuples (x, y, size, type) 障碍物信息，假设类型为矩形，轴对齐
        step_size: 计算斥力时移动的步长
        k_drone: 无人机的斥力系数
        k_obstacle: 障碍物的斥力系数
    返回:
        最佳逃离方向的角度（弧度），0表示正东方向，逆时针增加
    """
    x0=escape_drone.x
    y0=escape_drone.y       
    id0=escape_drone.id
    num_directions = 360
    best_direction_rad = 0
    min_repulsion = float('inf')
    
    for angle_deg in range(0, num_directions, 5):
        angle_rad = math.radians(angle_deg)
        x_new = x0 + step_size * math.cos(angle_rad)
        y_new = y0 + step_size * math.sin(angle_rad)
        total_repulsion = 0.0
        
        for drone in other_drones:
            x_d, y_d, id_d = drone.x, drone.y, drone.id
            dist = math.sqrt((x_new - x_d)**2 + (y_new - y_d)**2)
            if dist < 1e-5:
                repulsion = 1e10
            else:
                repulsion = k_drone / (dist**2)
            total_repulsion += repulsion
        
        for obs in obstacles:
            x_obs, y_obs, size_obs, type_obs = obs.x, obs.y, obs.size, obs.shape
            half_size = size_obs / 2.0
            x_min = x_obs - half_size
            x_max = x_obs + half_size
            y_min = y_obs - half_size
            y_max = y_obs + half_size
            
            if x_new < x_min:
                dx = x_min - x_new
            elif x_new > x_max:
                dx = x_new - x_max
            else:
                dx = 0.0
            if y_new < y_min:
                dy = y_min - y_new
            elif y_new > y_max:
                dy = y_new - y_max
            else:
                dy = 0.0
            dist_obs = math.sqrt(dx*dx + dy*dy)
            if dist_obs < 1e-5:
                repulsion_obs = 1e10
            else:
                repulsion_obs = k_obstacle / (dist_obs**2)
            total_repulsion += repulsion_obs
        
        if total_repulsion < min_repulsion:
            min_repulsion = total_repulsion
            best_direction_rad = angle_rad
    
    return best_direction_rad

# 初始化场景
drones, obstacles = setup_scenario()
escape_sequence = []
iteration = 0

# 创建图形
fig, ax = plt.subplots(figsize=(10, 10))
ax.set_xlim(5, GRID_SIZE)
ax.set_ylim(5, 25)
ax.set_aspect('equal')
ax.grid(True)

# 绘制障碍物
obstacle_patches = []
for obstacle in obstacles:
    if obstacle.shape == 'circle':
        patch = Circle((obstacle.x, obstacle.y), obstacle.size, fill=True, color='red', alpha=0.5)
    else:
        patch = Rectangle((obstacle.x - obstacle.size, obstacle.y - obstacle.size), 
                         2 * obstacle.size, 2 * obstacle.size, fill=True, color='red', alpha=0.5)
    obstacle_patches.append(patch)
ax.add_collection(PatchCollection(obstacle_patches, match_original=True))

# 绘制安全距离圆圈
safety_circles = []
for drone in drones:
    circle = Circle((drone.x, drone.y), SAFE_DISTANCE, fill=False, color='blue', linestyle='--', alpha=0.3)
    safety_circles.append(circle)
ax.add_collection(PatchCollection(safety_circles, match_original=True))

# 绘制无人机
drone_scatter = ax.scatter([d.x for d in drones], [d.y for d in drones], c='blue', s=100)
selected_drone_scatter = ax.scatter([], [], c='green', s=150)
path_line, = ax.plot([], [], 'g--', alpha=0.5)

# 添加文本显示
text = ax.text(GRID_SIZE/2, GRID_SIZE-2, "", ha='center', va='center', fontsize=12, 
              bbox=dict(boxstyle="round,pad=0.5", facecolor="white", alpha=0.7))

# 检查所有无人机是否安全
def all_drones_safe(drones, safe_distance):
    for i, drone1 in enumerate(drones):
        for j, drone2 in enumerate(drones):
            if i != j and drone1.distance_to(drone2) < safe_distance:
                return False
        for obstacle in obstacles:
            if obstacle.distance_to(drone1.x, drone1.y) < safe_distance:
                return False
    return True

# 更新函数
def update(frame):
    global drones, escape_sequence, iteration
    
    # iteration += 1
    # if iteration > MAX_ITERATIONS:
    #     text.set_text("Reached maximum iterations!")
    #     return drone_scatter, selected_drone_scatter, path_line, text
    
    if all_drones_safe(drones, SAFE_DISTANCE):
        text.set_text("All drones are safe!")
        return drone_scatter, selected_drone_scatter, path_line, text

    
    selected_drone = find_outermost_drone(drones)
    drones.remove(selected_drone)
    #safety_circles.remove(safety_circles[selected_drone.id - 1])
    step_size = 1.0
    best_angle_rad = calculate_avoidance_direction(selected_drone, drones, obstacles,step_size=step_size, k_drone=1.0, k_obstacle=1.0)
    best_angle_deg = math.degrees(best_angle_rad)


    
    # 记录逃离顺序和方向
    direction_deg = best_angle_deg
    escape_sequence.append((selected_drone.id, direction_deg, step_size))
    
    # 将角度转换为弧度
    direction_rad = np.deg2rad(direction_deg)

    # 计算无人机新的位置
    dx = np.cos(direction_rad) * step_size  # x方向上的步长
    dy = np.sin(direction_rad) * step_size  # y方向上的步长

    # 更新无人机的坐标
    selected_drone.x += dx
    selected_drone.y += dy
    selected_drone.escaped = True
    
    
    # 更新选中的无人机
    selected_drone_scatter.set_offsets([(selected_drone.x, selected_drone.y)])
    
    # 绘制移动路径
    path_line.set_data([selected_drone.original_x, selected_drone.x], 
                      [selected_drone.original_y, selected_drone.y])
    
    return drone_scatter, selected_drone_scatter, path_line, text

# 创建动画
ani = animation.FuncAnimation(fig, update, frames=MAX_ITERATIONS, interval=1000, repeat=False)

plt.title("Drone Escape Direction Selection Algorithm with Safe Path Planning")
plt.xlabel("X Coordinate")
plt.ylabel("Y Coordinate")
plt.tight_layout()
plt.show()

# 输出逃离顺序
print("\nDrone escape sequence and directions:")
for i, (drone_id, direction, distance) in enumerate(escape_sequence, 1):
    print(f"{i}. Drone {drone_id}: {direction:.2f} degrees, distance: {distance:.2f}")